import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import matplotlib as mpl

# 设置 matplotlib 支持中文
mpl.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体，根据系统情况可替换其他中文字体
mpl.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# *********************************************************
# 创建一个批量归一化层
# 首先，计算均值和方差：在每个通道上，对一个批次内所有样本的特征值计算均值和方差
# 比如，对于一个批次内的 1 个样本，每个样本有 3 个通道，每个通道有 4x4 的特征图
# 那么，对于每个通道（比如先 out_channels[0] ），计算这 4 个样本在该通道上的均值和方差
# 然后，对其他通道重复上述操作
# *********************************************************

# 定义输入数据的通道数
out_channels = 3

# 创建一个批量归一化层
batch_norm = nn.BatchNorm2d(out_channels)

# 生成一个随机输入张量，形状为 (batch_size, channels, height, width)
input_tensor = torch.randn(1, out_channels, 4, 4)

# 进行批量归一化操作
output_tensor = batch_norm(input_tensor)

# 从归一化的通道维度（维度 1）计算总和、均值和方差
sum_channel_dim = output_tensor.sum(dim=1)
mean_channel_dim = output_tensor.mean(dim=1)
var_channel_dim = output_tensor.var(dim=1)

# print("输入张量:", input_tensor)
print("输出张量:", output_tensor)
print("通道维度（维度 1）的和形状:", sum_channel_dim.shape)
print("通道维度（维度 1）的和:", sum_channel_dim)
print("通道维度（维度 1）的均值形状:", mean_channel_dim.shape)
print("通道维度（维度 1）的均值:", mean_channel_dim)
print("通道维度（维度 1）的方差形状:", var_channel_dim.shape)
print("通道维度（维度 1）的方差:", var_channel_dim)

# 获取第一个样本的数据
first_sample_sum = sum_channel_dim[0].flatten().detach().numpy()
first_sample_mean = mean_channel_dim[0].flatten().detach().numpy()
first_sample_var = var_channel_dim[0].flatten().detach().numpy()

# 生成 x 轴坐标
x = range(len(first_sample_sum))

# 绘制柱状图
plt.figure(figsize=(15, 5))
plt.subplot(1, 2, 1)
bar_width = 0.2
plt.bar([i - bar_width for i in x], first_sample_sum, width=bar_width, label='通道维度求和')
plt.bar(x, first_sample_mean, width=bar_width, label='通道维度均值')
plt.bar([i + bar_width for i in x], first_sample_var, width=bar_width, label='通道维度方差')
plt.title('柱状图展示通道维度统计量差异')
plt.xlabel('位置')
plt.ylabel('值')
plt.legend()

# 绘制折线图
plt.subplot(1, 2, 2)
plt.plot(x, first_sample_sum, marker='o', label='通道维度求和')
plt.plot(x, first_sample_mean, marker='s', label='通道维度均值')
plt.plot(x, first_sample_var, marker='^', label='通道维度方差')
plt.title('折线图展示通道维度统计量差异')
plt.xlabel('位置')
plt.ylabel('值')
plt.legend()

plt.tight_layout()
plt.show()

print("通道维度（维度 1）的均值形状:", mean_channel_dim.shape)
print("通道维度（维度 1）的均值:", mean_channel_dim)
print("通道维度（维度 1）的方差形状:", var_channel_dim.shape)
print("通道维度（维度 1）的方差:", var_channel_dim)
print("通道维度（维度 1）的和:", sum_channel_dim)

print("输入张量形状:", input_tensor.shape)
print("输出张量形状:", output_tensor.shape)